import numpy as np
import matplotlib.pyplot as plt
import matplotlib
from hera_cal import io, redcal, apply_cal, abscal, utils
from hera_cal.smooth_cal import build_time_blacklist
from hera_qm.metrics_io import load_metric_file
import glob
from hera_cal.datacontainer import DataContainer
from pyuvdata import UVData
from tqdm.notebook import tqdm
from matplotlib import colors
%matplotlib notebook

JD = 2459139
files = sorted(glob.glob('/lustre/aoc/projects/hera/H4C/2459139/zen.2459139.[3-4]*.sum.uvh5'))
ants = [(an, 'Jnn') for an in [24, 86, 117, 121, 138, 167, # Bad ants
                               25, 85, 99, 102, 118, 119, 120, 147]] # Good ants
bls = []
for n, a1 in enumerate(ants):
    for a2 in ants[n:]:
        bls.append(utils.join_bl(a1, a2))
        
jumps = {24: 927, # maps antenna number to integration
         86: 718,
         117: 974,
         121: 966,
         138: 1133,
         167: 150}
bad_to_good = {24: 25, # makes 14 m baselines
               86: 85,
               117: 118,
               121: 120,
               138: 118,
               167: 147}
bad_to_good_bl = {24: (119, 120, 'nn'),
                  86: (120, 119, 'nn'),
                  117: (119, 120, 'nn'),
                  121: (120, 119, 'nn'),
                  138: (118, 99, 'nn'),
                  167: (120, 102, 'nn')}

data_dict = {}
for f in tqdm(files):
    hd = io.HERAData(f)
    data, _, _ = hd.read(bls=bls)
    data_dict[f] = data

data = DataContainer({bl: np.vstack([data_dict[f][bl] for f in files]) for bl in bls})

times = np.hstack([data_dict[f].times for f in files])

# for ant in ants:
#     if ant[0] in jumps:
#         bl = utils.join_bl(ant, ant)
#         plt.figure()
#         to_plot = np.abs(data[bl])[(jumps[ant[0]]-10):(jumps[ant[0]]+11), :]
#         plt.imshow(to_plot, aspect='auto', norm=colors.LogNorm(np.median(to_plot)/2, np.median(to_plot)*2))
#         plt.colorbar()
#         plt.title(bl)

for ant in [a for a in ants if a[0] in jumps]:
    bad_auto = utils.join_bl(ant, ant)
    good_ant = (bad_to_good[ant[0]], ant[1])
    good_auto = utils.join_bl(good_ant, good_ant)
    cross = utils.join_bl(ant, good_ant)
    good_red = bad_to_good_bl[ant[0]]
    jump = jumps[ant[0]]
    
    titles = {bad_auto: 'Jumpy Auto',
              good_auto: 'Good Auto',
              cross: 'Jumpy-Good Cross',
              good_red: 'Redundant Good Cross'}
    
    print(f'Now examining the discontinuity in {ant} at JD = {times[jump]}. {good_ant} is a nearby good antenna. {good_red} is redundant with {cross}.')
    
    # Plot amplitude waterfalls
    plt.figure(figsize=(16,5))
    for bl, sp, in zip([bad_auto, good_auto, cross, good_red], [141, 142, 143, 144]):
        plt.subplot(sp)
        to_plot = np.abs(data[bl][jump-10:jump+11,:])
        plt.imshow(to_plot, norm=colors.LogNorm(np.median(to_plot)/2, np.median(to_plot)*2), aspect='auto',
                   extent=[hd.freqs[0]/1e6, hd.freqs[-1]/1e6, times[jump+10]-JD, times[jump-10]-JD])
        plt.ylabel(f'JD - {JD}')
        plt.xlabel('Frequency (MHz)')
        plt.colorbar()
        plt.title(f'{titles[bl]}: {bl}\nVisibility Amplitude')    
    plt.tight_layout()
    
    # Plot mean spectra before and after jump    
    plt.figure(figsize=(16,5))
    for bl, sp, in zip([bad_auto, good_auto, cross, good_red], [141, 142, 143, 144]):
        plt.subplot(sp)
   
        plt.semilogy(hd.freqs/1e6, np.mean(np.abs(data[bl][jump-10:jump,:]), axis=0), 
                     label=f'10 Integrations Before Jump:\n {np.around(np.mean(times[jump-10:jump]),5)}')
        plt.semilogy(hd.freqs/1e6, np.mean(np.abs(data[bl][jump+1:jump+11,:]), axis=0), 
                     label=f'10 Integrations After Jump:\n {np.around(np.mean(times[jump+1:jump+11]),5)}')
        plt.legend()
        plt.title(f'{titles[bl]}: {bl}')
        plt.xlabel('Frequency (MHz)')
        plt.ylabel('Uncalibrated Visibility Amplitude')
    plt.tight_layout()

    # Plot amplitude ratios    
    plt.figure(figsize=(16,5))
    for bl in [good_red, cross, good_auto, bad_auto]:
        plt.plot(hd.freqs/1e6, np.mean(np.abs(data[bl][jump+1:jump+11,:]), axis=0) / np.mean(np.abs(data[bl][jump-10:jump,:]), axis=0), label=f'{titles[bl]}: {bl}')
    plt.ylim([.5, 2])
    plt.plot(hd.freqs/1e6, 
             np.sqrt(np.mean(np.abs(data[bad_auto][jump+1:jump+11,:]), axis=0) / np.mean(np.abs(data[bad_auto][jump-10:jump,:]), axis=0) * \
                     np.mean(np.abs(data[good_auto][jump+1:jump+11,:]), axis=0) / np.mean(np.abs(data[good_auto][jump-10:jump,:]), axis=0)), label='Geometric Mean of Autos')
    plt.legend()
    plt.ylabel('Visibility Amplitude Ratio: After Jump / Before Jump ')
    plt.xlabel('Frequency (MHz)')
    plt.title('Average Amplitude of 10 Integrations, Ratio After / Before Jump')
    plt.tight_layout()

    # Plot phases before and after jump    
    plt.figure(figsize=(16,5))
    plt.plot(hd.freqs/1e6, np.angle(np.mean(data[cross][jump-20:jump-10,:], axis=0)), 
             label=f'{titles[cross]} {cross} Previous 10 Integrations:\n {np.around(np.mean(times[jump-20:jump-10]),5)}')
    plt.plot(hd.freqs/1e6, np.angle(np.mean(data[cross][jump-10:jump,:], axis=0)), 
             label=f'{titles[cross]} {cross} 10 Integrations Before Jump:\n {np.around(np.mean(times[jump-10:jump]),5)}')
    plt.plot(hd.freqs/1e6, np.angle(np.mean(data[cross][jump+1:jump+11,:], axis=0)), 
             label=f'{titles[cross]} {cross} 10 Integrations After Jump:\n {np.around(np.mean(times[jump+1:jump+11]),5)}')
    plt.plot(hd.freqs/1e6, np.angle(np.mean(data[cross][jump+11:jump+21,:], axis=0)), 
             label=f'{titles[cross]} {cross} Next 10 Integrations:\n {np.around(np.mean(times[jump+11:jump+21]),5)}')
    plt.ylabel('Phases of Mean Visibilities')
    plt.xlabel('Frequency (MHz)')
    plt.legend()
    plt.title('Phase of the Mean Visibility For Four Sets of 10 Integrations, Two Before and Two After the Jump')
    plt.tight_layout()
    
    # Plot phase differences before and after jump
    plt.figure(figsize=(16,5))
    for bl, lw in zip([cross, good_red], [2,.5]):
        angles = [np.angle(np.mean(data[bl][jump-20:jump-10,:], axis=0)),
                  np.angle(np.mean(data[bl][jump-10:jump,:], axis=0)),
                  np.angle(np.mean(data[bl][jump+1:jump+11,:], axis=0)),
                  np.angle(np.mean(data[bl][jump+11:jump+21,:], axis=0))]
        ang_diffs = np.diff(angles, axis=0)
        ang_diffs[ang_diffs < -np.pi] += 2*np.pi
        ang_diffs[ang_diffs > np.pi] -= 2*np.pi
        for ang_diff, label in zip(ang_diffs, ['Diff Before Jump', 'Diff Over Jump', 'Diff After Jump']):
            plt.plot(hd.freqs/1e6, ang_diff, lw=lw, label=f'{titles[bl]} {bl} {label}')
    plt.legend()
    plt.ylabel('Difference in Phases of Mean Visibilities')
    plt.xlabel('Frequency (MHz)')
    plt.tight_layout()

Now examining the discontinuity in (24, 'Jnn') at JD = 2459139.403693439. (25, 'Jnn') is a nearby good antenna. (119, 120, 'nn') is redundant with (24, 25, 'nn').

Now examining the discontinuity in (86, 'Jnn') at JD = 2459139.3803171846. (85, 'Jnn') is a nearby good antenna. (120, 119, 'nn') is redundant with (86, 85, 'nn').

Now examining the discontinuity in (117, 'Jnn') at JD = 2459139.4089503. (118, 'Jnn') is a nearby good antenna. (119, 120, 'nn') is redundant with (117, 118, 'nn').

Now examining the discontinuity in (121, 'Jnn') at JD = 2459139.408055515. (120, 'Jnn') is a nearby good antenna. (120, 119, 'nn') is redundant with (121, 120, 'nn').

Now examining the discontinuity in (138, 'Jnn') at JD = 2459139.426734149. (118, 'Jnn') is a nearby good antenna. (118, 99, 'nn') is redundant with (138, 118, 'nn').

More than 20 figures have been opened. Figures created through the pyplot interface (`matplotlib.pyplot.figure`) are retained until explicitly closed and may consume too much memory. (To control this warning, see the rcParam `figure.max_open_warning`).

More than 20 figures have been opened. Figures created through the pyplot interface (`matplotlib.pyplot.figure`) are retained until explicitly closed and may consume too much memory. (To control this warning, see the rcParam `figure.max_open_warning`).

More than 20 figures have been opened. Figures created through the pyplot interface (`matplotlib.pyplot.figure`) are retained until explicitly closed and may consume too much memory. (To control this warning, see the rcParam `figure.max_open_warning`).

More than 20 figures have been opened. Figures created through the pyplot interface (`matplotlib.pyplot.figure`) are retained until explicitly closed and may consume too much memory. (To control this warning, see the rcParam `figure.max_open_warning`).

More than 20 figures have been opened. Figures created through the pyplot interface (`matplotlib.pyplot.figure`) are retained until explicitly closed and may consume too much memory. (To control this warning, see the rcParam `figure.max_open_warning`).

Old code¶

# JD = 2459117
# files = sorted(glob.glob('/lustre/aoc/projects/hera/H4C/2459117/*.sum.uvh5'))
# bls = [(45,45,'ee'), (103, 103, 'ee'), (45, 103, 'ee'), (44, 44, 'ee'), (44, 45, 'ee'), (44, 103, 'ee')]

# plt.figure(figsize=(15,8))
# for bl, sp, vlim, in zip(bls[0:3], [131, 132, 133], [(1e6, 1e8), (1e6, 1e8), (1e4, 1e6)]):
#     plt.subplot(sp)
#     plt.imshow(np.abs(data[bl]), aspect='auto', norm=colors.LogNorm(*vlim),
#                extent=[hd.freqs[0]/1e6, hd.freqs[-1]/1e6, times[-1]-JD, times[0]-JD])
#     plt.colorbar(label='Amplitude (Uncalibrated)')
#     plt.ylabel('JD - 2459117')
#     plt.xlabel('Frequency (MHz)')
# #     plt.ylim([.25, .3])
#     plt.title(bl)
# plt.tight_layout()

# plt.figure(figsize=(12,5))
# for bl, sp, vlim, in zip([(45,45,'ee'), (44, 44, 'ee'), (44, 45, 'ee')],
#                          [131, 132, 133], [(4e6, 1e8), (1e6, 1e8), (1e4, 1e6)]):
#     plt.subplot(sp)
#     plt.imshow(np.abs(data[bl][128:149,:]), norm=colors.LogNorm(*vlim), aspect='auto',
#                extent=[hd.freqs[0]/1e6, hd.freqs[-1]/1e6, times[149]-2459117, times[128]-2459117])
#     plt.ylabel('JD - 2459117')
#     plt.xlabel('Frequency (MHz)')
#     plt.colorbar()
#     plt.title(f'{bl} Visibility Amplitude')    
# plt.tight_layout()
    
# plt.figure(figsize=(12,5))
# for bl, sp, vlim, in zip([(45,45,'ee'), (44, 44, 'ee'), (44, 45, 'ee')],
#                          [131, 132, 133], [(4e6, 1e8), (1e6, 1e8), (1e4, 1e6)]):
#     plt.subplot(sp)
#     plt.semilogy(hd.freqs/1e6, np.mean(np.abs(data[bl][128:138,:]), axis=0), 
#                  label=f'10 Integrations Before Jump:\n {np.around(np.mean(times[128:138]),5)}')
#     plt.semilogy(hd.freqs/1e6, np.mean(np.abs(data[bl][139:149,:]), axis=0), 
#                  label=f'10 Integrations After Jump:\n {np.around(np.mean(times[139:149]),5)}')
#     plt.legend()
#     plt.ylim(vlim)
#     plt.title(bl)
#     plt.xlabel('Frequency (MHz)')
#     plt.ylabel('Uncalibrated Visibility Amplitude')
# plt.tight_layout()

# plt.figure(figsize=(12,5))
# for bl in [(45,45,'ee'), (44, 44, 'ee'), (44, 45, 'ee')]:
#     plt.plot(hd.freqs/1e6, np.mean(np.abs(data[bl][139:149,:]), axis=0) / np.mean(np.abs(data[bl][128:138,:]), axis=0), label=bl)
# plt.ylim([.5, 2])
# plt.plot(hd.freqs/1e6, 
#          np.sqrt(np.mean(np.abs(data[(45,45,'ee')][139:149,:]), axis=0) / np.mean(np.abs(data[(45,45,'ee')][128:138,:]), axis=0) * \
#                  np.mean(np.abs(data[(44,44,'ee')][139:149,:]), axis=0) / np.mean(np.abs(data[(44,44,'ee')][128:138,:]), axis=0)), label='44 * 45 Geometric Mean')
    
# plt.legend()
# plt.ylabel('Visibility Amplitude Ratio: After Jump / Before Jump ')
# plt.xlabel('Frequency (MHz)')
# plt.title('Average Amplitude of 10 Integrations, Ratio After / Before Jump')

# plt.figure(figsize=(12,5))
# plt.plot(hd.freqs/1e6, np.mean(np.angle(data[bl][118:128,:]), axis=0), 
#          label=f'Previous 10 Integrations:\n {np.around(np.mean(times[118:128]),5)}')
# plt.plot(hd.freqs/1e6, np.mean(np.angle(data[bl][128:138,:]), axis=0), 
#          label=f'10 Integrations Before Jump:\n {np.around(np.mean(times[128:138]),5)}')
# plt.plot(hd.freqs/1e6, np.mean(np.angle(data[bl][139:149,:]), axis=0), 
#          label=f'10 Integrations After Jump:\n {np.around(np.mean(times[139:149]),5)}')
# plt.plot(hd.freqs/1e6, np.mean(np.angle(data[bl][149:159,:]), axis=0), 
#          label=f'Next 10 Integrations:\n {np.around(np.mean(times[149:159]),5)}')

# plt.ylabel('Mean Phase')
# plt.xlabel('Frequency (MHz)')
# plt.legend()
# plt.title('Mean Phase For Four Sets of 10 Integrations, Two Before and Two After the Jump')

# plt.figure()
# plt.plot(hd.freqs/1e6, np.mean(np.angle(data[bl][128:138,:]), axis=0), 
#              label=f'10 Integrations Before Jump:\n {np.around(np.mean(times[128:138]),5)}')
# plt.plot(hd.freqs/1e6, np.mean(np.angle(data[bl][139:149,:]), axis=0), 
#              label=f'10 Integrations After Jump:\n {np.around(np.mean(times[139:149]),5)}')

# for bl in bls:
#     plt.figure(figsize=(10,6))
#     plt.semilogy(times - 2459117, data[bl][:,10],  '.-', lw=.5, ms=.5, label='10: no RFI, low band')
#     plt.semilogy(times - 2459117, data[bl][:,400], '.-', lw=.5, ms=.5, label='400: FM')
#     plt.semilogy(times - 2459117, data[bl][:,485], '.-', lw=.5, ms=.5, label='485: FM')
#     plt.semilogy(times - 2459117, data[bl][:,739], '.-', lw=.5, ms=.5, label='739: ORBCOMM')
#     plt.semilogy(times - 2459117, data[bl][:,744], '.-', lw=.5, ms=.5, label='744: ORBCOMM')
#     plt.semilogy(times - 2459117, data[bl][:,750], '.-', lw=.5, ms=.5, label='750: no RFI')
#     plt.semilogy(times - 2459117, data[bl][:,1340],'.-', lw=.5, ms=.5,  label='1340: no RFI, high band')
#     plt.semilogy(times - 2459117, data[bl][:,1445],'.-', lw=.5, ms=.5,  label='1445: TV')
#     plt.legend()
#     plt.xlabel('JD - 2459117')
#     plt.ylabel(f'Raw Visibility Amplitude')
#     plt.title(f'{bl}')
#     plt.ylim([1e6, 2e8])